#include <stdio.h>
#include <mpi.h>
#include <math.h>
#include <stdlib.h>

//#define PRINT_TO_TERMINAL	
//PROC_CNT is defined by -D option in Makefile
#define	SLAVE_CNT					(PROC_CNT - 1) 
// N and K are defined in Makefile
//#define N							8	//queen size
//#define K							1	//depth K
#define MAX_SOLUTIONS				(1000000)
#define SOLUTION_FILE				"./solutions.txt"
#define	SOLUTION_FILE_SIZE			1000000
#define MAX_RECV_SIZE				10000

int row_pos[N]; //each element stors the position of queen of each column
int local_solution[MAX_SOLUTIONS][N];
int solution[MAX_SOLUTIONS][N];
int local_solution_cnt = 0; //the total number of solutions
int solution_cnt = 0; //the total number of solutions
int tag = 4242; //0-32767 can be used as tags

struct waiting_lst_St {
	int node[SLAVE_CNT];
	int *p_start;
	int *p_end;
	int full;
	int empty;
	int cnt;
	int diff; // the difference between pop times and push times
} waiting_lst;

// check if the position in the column meet the requirements
int attack_check(int *addr, int column)
{
	int i;
	for(i = 0; i < column; i++) {
		if ((addr[i] == addr[column]) || \
			(abs(addr[i] -addr[column]) == column - i)) {
			return 0;
		}	
	}		
	return 1;
}

void init_row_pos(void)
{
	int i;
	for(i = 0; i < N; i++)
		row_pos[i] = -1;
}

char txt_content[SOLUTION_FILE_SIZE];
void print_solution(void)
{
	int i, j;
	FILE *fp;
	int cnt = 0;

#ifdef PRINT_TO_TERMINAL	
	printf("Found %d solutions for %d-Queens Problem:\r\n", solution_cnt, N);
#endif
	cnt += sprintf(txt_content + cnt, "Found %d solutions for %d-Queens Problem:\r\n", solution_cnt, N);

	for (i = 0; i < solution_cnt; i++) {
		
		cnt += sprintf(txt_content + cnt, "S%d:  ", i + 1);
		for (j = 0; j < N; j++) {
#ifdef PRINT_TO_TERMINAL	
			printf("%4d", solution[i][j] + 1);
#endif
			cnt += sprintf(txt_content + cnt,"%4d", solution[i][j] + 1);
		}	
#ifdef PRINT_TO_TERMINAL	
		printf("\r\n");
#endif
		cnt += sprintf(txt_content + cnt,"\r\n");
	}
	
	fp = fopen("original_seq.txt", "wb");
	fwrite(txt_content, sizeof(char), cnt, fp);
	fclose(fp);
	
	
}

int init_waiting_lst()
{
	int i;
	for (i = 0; i < SLAVE_CNT; i++) {
		waiting_lst.node[i] = i + 1;
	}
	waiting_lst.p_start =  &waiting_lst.node[0];
	waiting_lst.p_end =  &waiting_lst.node[SLAVE_CNT - 1];
	waiting_lst.full = 1;
	waiting_lst.empty = 0;
	waiting_lst.cnt = SLAVE_CNT; // all slavers are free at the begining
	waiting_lst.diff = 0;
	return 0;
}

int pop_waiting_lst()
{
	int *bkp_start = waiting_lst.p_start;

	if (waiting_lst.cnt > 0) { // not empty, move pointers
		
		if (waiting_lst.p_start == &waiting_lst.node[SLAVE_CNT - 1]) {
		// reach the end of the buffer
			waiting_lst.p_start = &waiting_lst.node[0];
		} else {
			waiting_lst.p_start++;	
		}
		
		waiting_lst.cnt--;	
		waiting_lst.diff++;

		if (waiting_lst.cnt == 0) {
			//the fifo is empty
			waiting_lst.p_end = waiting_lst.p_start;
			waiting_lst.empty = 1;
			waiting_lst.full = 0;;
		}
		return *bkp_start;
	} else {
		return -1;
	}
}

int push_waiting_lst( int rank)
{
	if (waiting_lst.cnt < SLAVE_CNT) { // not full, move pointers
		if (waiting_lst.cnt != 0) { // move pointers
			if (waiting_lst.p_end == &waiting_lst.node[SLAVE_CNT - 1]) {
			// reach the end of the buffer
				waiting_lst.p_end = &waiting_lst.node[0];
			} else {
					waiting_lst.p_end++;	
			}
		}
		
		*waiting_lst.p_end = rank;
		waiting_lst.cnt++;	
		waiting_lst.diff--;

		if (waiting_lst.cnt == SLAVE_CNT) {
			//the fifo is full
			waiting_lst.empty = 0;
			waiting_lst.full = 1;
		}
		return 0;
	} else {
		return -1;
	}
}

// check waiting list is empty or not
int waiting_lst_empty(void)
{
	if (waiting_lst.empty)
		return 1;		
	else 
		return 0;
}

int print_waiting_lst(void)
{
	int i;
	int *p = waiting_lst.p_start;
	if (waiting_lst.empty) {
		printf("\r\nwaiting list empty\r\n");
		return 0;
	}
	printf("Waiting node list: %d nodes\r\n", waiting_lst.cnt);
	for(i = 0; i < waiting_lst.cnt; i++) {
		printf("%d  ", *p);
		p++;
		if (p > &waiting_lst.node[SLAVE_CNT - 1])
			p = &waiting_lst.node[0];
	}
	printf("\r\n");
	return 1;
}

int check_waiting_list(MPI_Request *recv_req, MPI_Status *status)
{
	static int round = 0;
	int recv_flag = 1;
	int len, err;

	if (round > 0) {
		if (waiting_lst_empty()) { 
			// waiting list is empty, then wait here
			MPI_Wait(recv_req, status); 
			recv_flag = 1;
		} else { 
				MPI_Test(recv_req, &recv_flag, status); 
		}
	}

	if (recv_flag) { 
		// Master node finished one non-blocking receiving
		if (round > 0) {
			push_waiting_lst((*status).MPI_SOURCE);
			MPI_Get_count(status, MPI_INT, &len);
			solution_cnt +=  len / N ;
		}
		err = MPI_Irecv(solution + solution_cnt, MAX_RECV_SIZE, MPI_INT,  \
			MPI_ANY_SOURCE, tag, MPI_COMM_WORLD, recv_req);
		round++;
	}
	return 0;
}

int n_queens(int rank)
{
	int i, len, err, dispatch = 0;
	int waiting_node_rank;
	int start_column, end_column, cur_column;

	start_column = 0;
	end_column = K - 1;
	cur_column = 0;
	MPI_Request send_req[PROC_CNT];
	MPI_Status status;
	MPI_Request recv_req;

	if (rank == 0) {
		init_waiting_lst(); // init the waiting list of slave nodes
	} 

	while (1) {     
		init_row_pos();
		if (rank != 0) {
			start_column = K;
			end_column = N -1;
			cur_column = K;

			//slave nodes wait for rank 0 to dispatch task here
			MPI_Recv(row_pos, K, MPI_INT, 0, tag, MPI_COMM_WORLD, &status);
			//if received_data[0] < 0, then slave stops and return
			if (row_pos[0] < 0)
				return 1;
		}

		while (cur_column >= start_column) {
			do { // increase row position until find one valid positon
				row_pos[cur_column]++;
				if (row_pos[cur_column] >= N)
					break;  
			} while (!attack_check(row_pos, cur_column));

			if (row_pos[cur_column] >= N) {
				// out of range, then roll back one column
				row_pos[cur_column] = -1;
				cur_column--;
			} else {
				if (cur_column == end_column) {
					//reach the final column, don't increase column
					//get one solution at each entering
					if (rank == 0) {
						check_waiting_list(&recv_req, &status); 
								/* 1. non-block receiving from slave
									nodes. Add these nodes to waiting_list when 
									finished receiving.
									2. Wait here for non-block receiving to finish 
									if waiting list is empty.
									*/
						waiting_node_rank = pop_waiting_lst();
						MPI_Isend(row_pos, K, MPI_INT, waiting_node_rank, tag, \
								MPI_COMM_WORLD, &send_req[waiting_node_rank]);
						dispatch++;
					} else { 
						// slave nodes : copy one solution to a solution array
						memcpy((void *)local_solution[local_solution_cnt], \
								(void *)row_pos, sizeof(row_pos));
						local_solution_cnt++;
					}
						
				} else { 
					// not teh last column, then increase column
					cur_column++;   
				}
			}
		}
		
		/* master node received all the solution from slave 
		  nodes of the first k columns, then return	
		*/
		if (rank == 0 && cur_column < 0) {
			/* Because rank 0 has no tasks to dispatch then, 
			   so we can use blocking send here.
			*/
			MPI_Wait(&recv_req, &status); //wait for the last non-blocking recv
			push_waiting_lst(status.MPI_SOURCE);
			MPI_Get_count(&status, MPI_INT, &len);
			solution_cnt +=  len / N ;
			while (waiting_lst.diff) { // wait for the left slave nodes
				err = MPI_Recv(solution + solution_cnt, MAX_RECV_SIZE, MPI_INT,  \
						MPI_ANY_SOURCE, tag, MPI_COMM_WORLD, &status);
				push_waiting_lst(status.MPI_SOURCE);
				MPI_Get_count(&status, MPI_INT, &len);
				solution_cnt +=  len / N ;
			}
			for (i = 1; i <= SLAVE_CNT; i++) { 
				//send stop signal to all the slavers
				MPI_Isend(&cur_column, 1, MPI_INT, i, tag, MPI_COMM_WORLD, \
							&send_req[waiting_node_rank]);   
			}
			return 1;
		}

		// slave nodes finished one task dispatched by master node
		if (rank != 0 && cur_column < start_column) {
			//slave node finished filling all the following columns, 
			//then send the whole solutions to master node
			MPI_Send(local_solution[0],  N * local_solution_cnt, MPI_INT, 0, \
										tag, MPI_COMM_WORLD);
			local_solution_cnt = 0;
		}

	} 
}


int main (int argc, char **argv)
 {
    int err, rank, size;
	double t0, t1;
    err = MPI_Init( &argc, &argv);
    err= MPI_Comm_rank( MPI_COMM_WORLD, &rank );
    err= MPI_Comm_size( MPI_COMM_WORLD, &size );
	
	t0 = MPI_Wtime();
	n_queens(rank); 		
	t1 = MPI_Wtime();

	if (rank == 0) {
		printf("Found %d solutions for %d-Queens Problem:\r\n", solution_cnt, N);
		printf("###Time assuming: %f us\r\n", (t1 - t0) * 1000000);
		printf("##pop push diff: %d\r\n", waiting_lst.diff);
		print_solution();
	} 
	printf("\r\n###rank %d stopped\r\n", rank);
	
    err = MPI_Finalize();
    return err;
}
